import os
import sys
import cv2
import numpy as np
from PyQt5.QtCore import QThread, pyqtSignal
import time
import multiprocessing
from concurrent.futures import ThreadPoolExecutor

class StereoGenerator(QThread):
    progress_updated = pyqtSignal(int)
    log_message = pyqtSignal(str)
    process_completed = pyqtSignal(bool, str)
    
    def __init__(self):
        super().__init__()
        self.is_running = False
        self.should_stop = False
        self.frames_dir = ""
        self.depth_dir = ""
        self.left_eye_dir = ""
        self.right_eye_dir = ""
        self.shift_strength = 0.05  # 默认视差强度，适合显示器/电视
        self.hardware_accel = "cpu"  # 硬件加速选项: "cpu", "cuda", "opencl"
        self.use_parallel = True  # 是否使用并行处理
        self.num_workers = max(1, multiprocessing.cpu_count() - 1)  # 工作线程数
    
    def run(self):
        self.is_running = True
        self.should_stop = False
        
        try:
            # 创建左右眼输出目录
            if not os.path.exists(self.left_eye_dir):
                os.makedirs(self.left_eye_dir)
            if not os.path.exists(self.right_eye_dir):
                os.makedirs(self.right_eye_dir)
            
            # 获取帧图像列表
            frame_files = sorted([f for f in os.listdir(self.frames_dir) if f.endswith('.jpg') or f.endswith('.png')])
            total_frames = len(frame_files)
            
            if total_frames == 0:
                self.log_message.emit("未找到帧图像，请先进行视频分帧")
                self.process_completed.emit(False, "未找到帧图像")
                self.is_running = False
                return
            
            self.log_message.emit(f"开始生成左右眼视图，共{total_frames}帧...")
            self.log_message.emit(f"硬件加速模式: {self.hardware_accel}")
            
            # 提取实际的硬件加速类型（去除描述文本）
            actual_accel = self.hardware_accel.split()[0] if " " in self.hardware_accel else self.hardware_accel
            
            # 检查CUDA可用性
            if "cuda" in actual_accel:
                if cv2.cuda.getCudaEnabledDeviceCount() > 0:
                    self.log_message.emit("CUDA加速可用")
                    self.log_message.emit(f"CUDA设备数: {cv2.cuda.getCudaEnabledDeviceCount()}")
                else:
                    self.log_message.emit("CUDA不可用，切换到CPU模式")
                    self.hardware_accel = "cpu"
                    actual_accel = "cpu"
            
            # 并行处理
            if self.use_parallel and total_frames > 10:
                self.log_message.emit(f"使用并行处理，工作线程数: {self.num_workers}")
                self.process_frames_parallel(frame_files, total_frames)
            else:
                # 处理每一帧
                self.process_frames_sequential(frame_files, total_frames)
            
            self.log_message.emit("左右眼视图生成完成")
            # 直接发送完成信号，不要在这里进行视频合成
            self.process_completed.emit(True, "左右眼视图生成完成")
            
        except Exception as e:
            import traceback
            self.log_message.emit(f"左右眼视图生成出错: {str(e)}")
            self.log_message.emit(f"详细错误: {traceback.format_exc()}")
            self.process_completed.emit(False, f"左右眼视图生成出错: {str(e)}")
        
        self.is_running = False
    
    def process_frames_sequential(self, frame_files, total_frames):
        """顺序处理帧"""
        for i, frame_file in enumerate(frame_files):
            if self.should_stop:
                self.log_message.emit("用户取消了立体视图生成")
                self.process_completed.emit(False, "用户取消")
                return
            
            self.process_single_frame(frame_file, i, total_frames)
            
            # 更新进度
            progress = int((i + 1) / total_frames * 100)
            self.progress_updated.emit(progress)
            
            if (i + 1) % 10 == 0 or (i + 1) == total_frames:
                self.log_message.emit(f"已处理 {i + 1}/{total_frames} 帧 ({progress}%)")
    
    def process_frames_parallel(self, frame_files, total_frames):
        """并行处理帧"""
        processed_count = 0
        
        # 创建线程池
        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
            # 提交所有任务
            future_to_frame = {executor.submit(self.process_single_frame, frame_file, i, total_frames): i 
                              for i, frame_file in enumerate(frame_files)}
            
            # 处理完成的任务
            for future in future_to_frame:
                if self.should_stop:
                    executor.shutdown(wait=False)
                    self.log_message.emit("用户取消了立体视图生成")
                    self.process_completed.emit(False, "用户取消")
                    return
                
                try:
                    future.result()  # 获取结果，如果有异常会抛出
                    processed_count += 1
                    
                    # 更新进度
                    progress = int(processed_count / total_frames * 100)
                    self.progress_updated.emit(progress)
                    
                    if processed_count % 10 == 0 or processed_count == total_frames:
                        self.log_message.emit(f"已处理 {processed_count}/{total_frames} 帧 ({progress}%)")
                        
                except Exception as e:
                    self.log_message.emit(f"处理帧时出错: {str(e)}")
    
    def process_single_frame(self, frame_file, index, total_frames):
        """处理单个帧"""
        # 读取原始帧
        frame_path = os.path.join(self.frames_dir, frame_file)
        frame = cv2.imread(frame_path)
        
        if frame is None:
            self.log_message.emit(f"无法读取图像: {frame_path}")
            return
        
        # 读取对应的深度图
        depth_filename = os.path.splitext(frame_file)[0] + '_depth.png'
        depth_path = os.path.join(self.depth_dir, depth_filename)
        
        if not os.path.exists(depth_path):
            self.log_message.emit(f"未找到深度图: {depth_path}")
            return
        
        depth_map = cv2.imread(depth_path, cv2.IMREAD_GRAYSCALE)
        
        # 生成左右眼视图
        left_eye, right_eye = self.generate_stereo_views(frame, depth_map)
        
        # 保存左右眼视图
        left_filename = os.path.splitext(frame_file)[0] + '.jpg'
        right_filename = os.path.splitext(frame_file)[0] + '.jpg'
        
        left_path = os.path.join(self.left_eye_dir, left_filename)
        right_path = os.path.join(self.right_eye_dir, right_filename)
        
        cv2.imwrite(left_path, left_eye)
        cv2.imwrite(right_path, right_eye)
    
    def generate_stereo_views(self, frame, depth_map):
        """生成左右眼视图，支持硬件加速"""
        # 获取图像尺寸
        height, width = frame.shape[:2]
        
        # 确保深度图与帧大小一致
        if depth_map.shape[:2] != frame.shape[:2]:
            depth_map = cv2.resize(depth_map, (width, height), interpolation=cv2.INTER_NEAREST)
        
        # 平滑深度图以减少噪声，同时保持边缘
        depth_map = cv2.bilateralFilter(depth_map, 9, 75, 75)
        
        # 归一化深度图 (0-1范围)
        if depth_map.max() > 1:
            depth_normalized = depth_map.astype(np.float32) / 255.0
        else:
            depth_normalized = depth_map.astype(np.float32)
        
        # 提取实际的硬件加速类型（去除描述文本）
        actual_accel = self.hardware_accel.split()[0] if " " in self.hardware_accel else self.hardware_accel
        
        # 使用CUDA加速
        if "cuda" in actual_accel and cv2.cuda.getCudaEnabledDeviceCount() > 0:
            return self.generate_stereo_views_cuda(frame, depth_normalized)
        
        # 使用向量化操作代替循环
        # 创建坐标网格
        x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height))
        
        # 计算视差位移，限制最大偏移量
        max_shift = width * 0.05  # 限制最大偏移为图像宽度的5%
        shift_map = (self.shift_strength * (1.0 - depth_normalized) * width)
        shift_map = np.clip(shift_map, 0, max_shift).astype(np.float32)
        
        # 创建映射矩阵
        map_left_x = (x_coords + shift_map / 2).astype(np.float32)
        map_left_y = y_coords.astype(np.float32)
        
        map_right_x = (x_coords - shift_map / 2).astype(np.float32)
        map_right_y = y_coords.astype(np.float32)
        
        # 使用重映射生成左右眼视图，使用双三次插值提高质量
        left_eye = cv2.remap(frame, map_left_x, map_left_y, cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
        right_eye = cv2.remap(frame, map_right_x, map_right_y, cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
        
        return left_eye, right_eye
    
    def generate_stereo_views_cuda(self, frame, depth_normalized):
        """使用CUDA加速生成左右眼视图"""
        try:
            # 获取图像尺寸
            height, width = frame.shape[:2]
            
            # 将图像上传到GPU
            gpu_frame = cv2.cuda_GpuMat()
            gpu_frame.upload(frame)
            
            # 创建坐标网格
            x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height))
            
            # 计算视差位移，限制最大偏移量
            max_shift = width * 0.05  # 限制最大偏移为图像宽度的5%
            shift_map = (self.shift_strength * (1.0 - depth_normalized) * width)
            shift_map = np.clip(shift_map, 0, max_shift).astype(np.float32)
            
            # 创建映射矩阵
            map_left_x = (x_coords + shift_map / 2).astype(np.float32)
            map_left_y = y_coords.astype(np.float32)
            
            map_right_x = (x_coords - shift_map / 2).astype(np.float32)
            map_right_y = y_coords.astype(np.float32)
            
            # 上传映射矩阵到GPU
            gpu_map_left_x = cv2.cuda_GpuMat()
            gpu_map_left_y = cv2.cuda_GpuMat()
            gpu_map_right_x = cv2.cuda_GpuMat()
            gpu_map_right_y = cv2.cuda_GpuMat()
            
            gpu_map_left_x.upload(map_left_x)
            gpu_map_left_y.upload(map_left_y)
            gpu_map_right_x.upload(map_right_x)
            gpu_map_right_y.upload(map_right_y)
            
            # 创建GPU输出矩阵
            gpu_left_eye = cv2.cuda_GpuMat()
            gpu_right_eye = cv2.cuda_GpuMat()
            
            # 使用CUDA重映射，使用双三次插值提高质量
            cv2.cuda.remap(gpu_frame, gpu_left_eye, gpu_map_left_x, gpu_map_left_y, cv2.INTER_CUBIC, cv2.BORDER_REPLICATE)
            cv2.cuda.remap(gpu_frame, gpu_right_eye, gpu_map_right_x, gpu_map_right_y, cv2.INTER_CUBIC, cv2.BORDER_REPLICATE)
            
            # 下载结果
            left_eye = gpu_left_eye.download()
            right_eye = gpu_right_eye.download()
            
            return left_eye, right_eye
            
        except Exception as e:
            self.log_message.emit(f"CUDA处理出错: {str(e)}，切换到CPU模式")
            # 回退到CPU处理
            return self.generate_stereo_views(frame, depth_normalized)
    
    def set_shift_strength(self, strength):
        """设置视差强度"""
        self.shift_strength = float(strength)
    
    def set_hardware_accel(self, accel_type):
        """设置硬件加速类型"""
        self.hardware_accel = accel_type
        self.log_message.emit(f"设置硬件加速为: {accel_type}")
        
        # 提取实际的硬件加速类型（去除描述文本）
        actual_accel = accel_type.split()[0] if " " in accel_type else accel_type
        
        # 检查CUDA可用性
        if "cuda" in actual_accel:
            if cv2.cuda.getCudaEnabledDeviceCount() > 0:
                self.log_message.emit(f"CUDA可用，设备数: {cv2.cuda.getCudaEnabledDeviceCount()}")
                self.log_message.emit(f"CUDA设备名称: {cv2.cuda.getDevice()}")
            else:
                self.log_message.emit("CUDA不可用，将在运行时切换到CPU模式")
    
    def set_parallel_processing(self, use_parallel, num_workers=None):
        """设置并行处理选项"""
        self.use_parallel = use_parallel
        if num_workers is not None:
            self.num_workers = num_workers
        self.log_message.emit(f"并行处理: {'启用' if use_parallel else '禁用'}, 工作线程数: {self.num_workers}")
    
    def stop(self):
        """停止处理"""
        self.should_stop = True